
import numpy as np
import pylab
import JT
import csv
import JT.JTconfig as config
import os
from JT import UploadToGoogle
import my_auth
from socket import gethostname
from datetime import datetime
import time

user = gethostname()

from win32api import GetSystemMetrics

#===============================================================================
#    Data Paths
#===============================================================================
fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
#fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'

roi = [[210,187],[277,370]]


class Test():
    def __init__(self,logging=True,upload=True):            

#===============================================================================
#    Google Data Upload Setup
#===============================================================================
        cred = my_auth.credentials()
        username = cred.name
        password = cred.pw
        self._connectSuccess = False
        
        if upload:
            try:
                self._uploader = UploadToGoogle.Uploader(username,password)
                # Choose newest sheet and book
                self._uploader._run(book='0',sheet='0')
                self._connectSuccess = True
            except:
                print 'Could not connect to Google'
        
        if logging:
            logTag = 'GD_Results'     #This tag can be used to identify trials
            logName = time.strftime('%Y%m%d_%H%M%S')+'_'+logTag+'.log'
            logDir = 'C:/Users/bryan/bryan-code/trunk/Results'
            print 'Final Result Logging dir: ', logDir 
            self._logFile = csv.writer(open(os.path.join(logDir,logName),'wt'),
                                   dialect='excel-tab')
            # Write Headers
            self._logFile.writerow(('Iteration','Cost','Starting Pose0(mm)',
                                    'Starting Pose1(mm)','Starting Pose2(mm)',
                                    'Starting Pose3(deg)','Starting Pose4(deg)',
                                    'Starting Pose5(deg)','Final Pose0(mm)',
                                    'Final Pose1(mm)','Final Pose2(mm)',
                                    'Final Pose3(deg)','Final Pose4(deg)',
                                    'Final Pose5(deg)'))
            
#===============================================================================
#    Setup DRR
#===============================================================================
        drr = JT.DRR()
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFileLat)
        cal.LoadStaFile(staFile)
        
        drr.SetExternalCalibration(cal)
        
        self.volume = JT.Volume()
        self.volume.SetVolumeFilename(inputVolFileName)
        self.volume.SetOriginalTransform(cal._VolumeT)
        self.volume.SetBodyCenter((3.906, 2.17, -7.378))
        #self.volume.UseRayCastMapper()
        drr.AddVolume(self.volume)
        
        drr._renWin.SetOffScreenRendering(1)
        
        # Set Fixed image using path
        self.fixedImage = JT.FixedImage()
        self.fixedImage.SetFileName(fixedImageName)

        self.reg=JT.Registration()
        self.reg.SetFixedImage(self.fixedImage.GetImage(0))
        self.reg.SetMovingImageGenerator(drr)
        self.reg.SetRegionOfInterest(roi)
        
        self.Testing = JT.Testing()


    def run_test(self,n_pose=50,level=4,method='gc',debug=True,plot=False,upload_data=True):    
    #===============================================================================
    #    Setup Optimizer
    #===============================================================================
    # Select the desired starting pose, between [0,200)
        
        # set Metric
        self.reg.SetImageMetric(method)
        self.reg._plot = True
        
        from JT.Optimizer import RegularStepGradientDescentOptimizer
        gd = RegularStepGradientDescentOptimizer()
        
        # Set the starting pose
        pose = self.Testing.GetStartPose(n_pose)
        gd.SetInitialPosition(pose)
        print "Starting Pose: ", pose
        
        # Set Scales 1/1000 for translations, 1 for rotations
        scales = [1.0,1.0,1.0,1000.0,1000.0,1000.0]
        gd.SetScales(scales)
        
        # Set the cost function
        gd.SetCostFunction(self.reg)
        
        # Modify Optimizer Parameters
        gd.SetMaximumStepLength(2)
        gd.SetMinimumStepLength(5e-3)
        gd.SetRelaxationFactor(0.8)
        gd.SetNumberOfIterations(50)
        
        # Start the optimizer
        gd._debug = False
        gd.StartOptimization()
        
        # Print Final State    
        gd.PrintSelf()
        
        # Save registration data to file for analysis of tuning parameters
        # Log status/results to a file that can be plotted later.
        initial_pose = gd._InitialPosition
        final_pose = gd._CurrentPosition
        final_cost = gd._Value
        dt = datetime.fromtimestamp(time.time())
        time_stamp = dt.strftime('%m/%d/%Y_%H:%M:%S')
        user = gethostname()
        
        my_tupel = [n_pose, final_cost, initial_pose[0], initial_pose[1],
                    initial_pose[2],np.degrees(initial_pose[3]), 
                    np.degrees(initial_pose[4]),np.degrees(initial_pose[5]),
                    final_pose[0],final_pose[1],final_pose[2],
                    np.degrees(final_pose[3]),np.degrees(final_pose[4]),
                    np.degrees(final_pose[5])]
        
        self._logFile.writerow(my_tupel)
        
        
        if upload_data and self._connectSuccess:
            output_dict = {}
            output_dict['user'] = user
            output_dict['time'] = time_stamp
            output_dict['iteration'] = str(n_pose)
            output_dict['cost'] = str(final_cost)
            output_dict['startingpose0'] = str(initial_pose[0])
            output_dict['startingpose1'] = str(initial_pose[1])
            output_dict['startingpose2'] = str(initial_pose[2])
            output_dict['startingpose3'] = str(np.degrees(initial_pose[3]))
            output_dict['startingpose4'] = str(np.degrees(initial_pose[4]))
            output_dict['startingpose5'] = str(np.degrees(initial_pose[5]))
            output_dict['finalpose0'] = str(final_pose[0])
            output_dict['finalpose1'] = str(final_pose[1])
            output_dict['finalpose2'] = str(final_pose[2])
            output_dict['finalpose3'] = str(np.degrees(final_pose[3]))
            output_dict['finalpose4'] = str(np.degrees(final_pose[4]))
            output_dict['finalpose5'] = str(np.degrees(final_pose[5]))
            
            self._uploader.Write(output_dict)
            #print "Data written to google spreadsheet!"
            #print record.content

if __name__ == "__main__":
    T = Test()
    for i in range(200):
        print "Pose: ", i
        T.run_test(method='bpc',n_pose=i,plot=False)
        # Reset the reg class for the next run
        T.reg.Reset()
        #JT.PlotResultsFromFile.run()
    
    #pylab.show()


